# -*- encoding: utf-8 -*-
"""
@File    :   Prepocess.py
@Contact :   thgpddl@163.com

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2021/9/29 22:20   thgpddl      1.0         None
"""
import numpy as np
import cv2


class PreProcess:
    def __init__(self, predict=True):
        self.predict = predict

    def __call__(self, img):
        """
        表情识别输入预处理
        :param img: 灰度图
        :return: 返回shape=[1,48,48,1]的array
        """
        if self.predict:
            assert len(img.shape) == 2, print("输入图像应当为灰度图：2通道")
            inputs = np.expand_dims(np.expand_dims(np.array(cv2.resize(img, (48, 48))).astype("float32"), axis=0),
                                    axis=-1)
            return inputs
        else:
            assert False, print("Preprocess参数eval应当为True")


class PostProcess:
    def __init__(self):
        self.label = ['anger', 'disgust', 'fear', 'happy', 'sad', 'surprised', 'normal']

    def __call__(self, res):
        index = np.argmax(res)

        conf = res[0, index]
        label = self.label[index]
        return label, conf
